﻿#include "VulkanMaterials.h"
#include "Render/PreRender/PreRenderMaterial.h"
#include "Render/PreRender/PreRenderTexture2D.h"
#include "Render/PreRender/PreRenderTextureCube.h"
#include "Render/PreRender/PreRenderShader.h"

#include "Render/Material/VulkanMaterial.h"
#include "Render/Material/VulkanShaders.h"
#include "Render/Material/VulkanShader.h"
#include "Render/Texture/VulkanTexture2Ds.h"
#include "Render/Texture/VulkanTextureCubes.h"
#include "Render/VulkanContext.h"


void CVulkanMaterials::UpdateMaterials(
    const TVector<std::shared_ptr<CPreRenderMaterial>>& InMaterials,
    CVulkanShaders* Shaders,
    CVulkanTexture2Ds* Texture2Ds,
    CVulkanTextureCubes* TextureCubes
    )
{
    Materials.Clear();

    for (auto& MaterialData : InMaterials)
    {
        auto PreRenderParameters = MaterialData->Parameters;

        if(!MaterialParameters.Contains(PreRenderParameters))
        {
            auto NewParameter = std::make_shared<CVulkanMaterialParameter>();
            MaterialParameters.Add(PreRenderParameters, NewParameter);

            NewParameter->FloatParameterBuffer.Resize(PreRenderParameters.Parameters.Num() * sizeof(float));
            if(PreRenderParameters.Parameters.Num()>0)
            {
                NewParameter->FloatParameterBuffer.WriteData(PreRenderParameters.Parameters.Data(), PreRenderParameters.Parameters.Num() * sizeof(float));
            }

            NewParameter->Texture2Ds.Reserve(PreRenderParameters.Texture2Ds.Num());
            for(auto& TextureData : PreRenderParameters.Texture2Ds)
            {
                NewParameter->Texture2Ds.Add(Texture2Ds->FindVulkanObj(TextureData->Pixels));
            }

            NewParameter->TextureCubes.Reserve(PreRenderParameters.TextureCubes.Num());
            for(auto& TextureData : PreRenderParameters.TextureCubes)
            {
                NewParameter->TextureCubes.Add(TextureCubes->FindVulkanObj(TextureData->WriteInData));
            }
        }
        
        auto Parameters=MaterialParameters[PreRenderParameters];

        auto PreRenderPipeline = MaterialData->Pipeline;
        if(!MaterialPipelines.Contains(PreRenderPipeline))
        {
            auto NewPipeline = std::make_shared<CVulkanPipeline>();
            MaterialPipelines.Add(PreRenderPipeline, NewPipeline);

            NewPipeline->Type = PreRenderPipeline.Type;

            //创建描述符布局
            {
                uint32_t CurrentBindingIndex = 0;
                TVector<vk::DescriptorSetLayoutBinding> DescriptorSetLayoutBindings;

                if(!Parameters->Texture2Ds.Empty())
                {
                    DescriptorSetLayoutBindings.Add(vk::DescriptorSetLayoutBinding()
                        .setBinding(CurrentBindingIndex)
                        .setDescriptorType(vk::DescriptorType::eCombinedImageSampler)
                        .setDescriptorCount(Parameters->Texture2Ds.Num())
                        .setStageFlags(vk::ShaderStageFlagBits::eFragment));
                    
                    CurrentBindingIndex++;         
                }

                if(!Parameters->TextureCubes.Empty())
                {
                    DescriptorSetLayoutBindings.Add(vk::DescriptorSetLayoutBinding()
                        .setBinding(CurrentBindingIndex)
                        .setDescriptorType(vk::DescriptorType::eCombinedImageSampler)
                        .setDescriptorCount(Parameters->TextureCubes.Num())
                        .setStageFlags(vk::ShaderStageFlagBits::eFragment));
                    
                    CurrentBindingIndex++;         
                }

                if(Parameters->FloatParameterBuffer.GetSize()>0)
                {
                    DescriptorSetLayoutBindings.Add(vk::DescriptorSetLayoutBinding()
                        .setBinding(CurrentBindingIndex)
                        .setDescriptorType(vk::DescriptorType::eUniformBuffer)
                        .setDescriptorCount(1)
                        .setStageFlags(vk::ShaderStageFlagBits::eVertex | vk::ShaderStageFlagBits::eFragment));
                    
                    CurrentBindingIndex++;
                }

                NewPipeline->DescriptorSetLayout = CVulkanContext::Get()->GetDevice().createDescriptorSetLayout(vk::DescriptorSetLayoutCreateInfo()
                    .setBindingCount(DescriptorSetLayoutBindings.Num())
                    .setPBindings(DescriptorSetLayoutBindings.Data()));
            }

            //创建管线布局
            {
                TVector<vk::DescriptorSetLayout> DescriptorSetLayouts =
                {
                    NewPipeline->DescriptorSetLayout,
                    CVulkanContext::Get()->GetSceneGlobal()->ModelDescriptorSetLayout.Layout, // model info
                    CVulkanContext::Get()->GetSceneGlobal()->SceneParameterLayout //场景全局数据
                };

                vk::PipelineLayoutCreateInfo PipelineLayoutCreateInfo = vk::PipelineLayoutCreateInfo()
                    .setSetLayouts(DescriptorSetLayouts)
                    ;

                NewPipeline->PipelineLayout = CVulkanContext::Get()->GetDevice().createPipelineLayout(PipelineLayoutCreateInfo);
            }

            //创建管线
            {
                TVector<vk::DynamicState> DynamicStates;
                DynamicStates.Add(vk::DynamicState::eViewport);
                DynamicStates.Add(vk::DynamicState::eScissor);

                vk::PipelineDynamicStateCreateInfo PipelineDynamicStateCreateInfo = vk::PipelineDynamicStateCreateInfo()
                    .setDynamicStates(DynamicStates);

                vk::PipelineShaderStageCreateInfo PipelineShaderStageCreateInfos[2];
                PipelineShaderStageCreateInfos[0] = vk::PipelineShaderStageCreateInfo()
                    .setStage(vk::ShaderStageFlagBits::eVertex)
                    .setModule(Shaders->FindVulkanObj(PreRenderPipeline.VertexShader->WriteInData)->Module)
                    .setPName("main");

                PipelineShaderStageCreateInfos[1] = vk::PipelineShaderStageCreateInfo()
                    .setStage(vk::ShaderStageFlagBits::eFragment)
                    .setModule(Shaders->FindVulkanObj(PreRenderPipeline.FragmentShader->WriteInData)->Module)
                    .setPName("main");

                TVector<vk::VertexInputBindingDescription> VertexInputBindingDescriptions=CVulkanContext::Get()->GetMeshVertexInputBindingDescriptions();
                TVector<vk::VertexInputAttributeDescription> VertexInputAttributeDescriptions = CVulkanContext::Get()->GetMeshVertexInputAttributeDescriptions();
                
                vk::PipelineVertexInputStateCreateInfo PipelineVertexInputStateCreateInfo = vk::PipelineVertexInputStateCreateInfo()
                    .setVertexBindingDescriptionCount(VertexInputBindingDescriptions.Num())
                    .setPVertexBindingDescriptions(VertexInputBindingDescriptions.Data())
                    .setVertexAttributeDescriptionCount(VertexInputAttributeDescriptions.Num())
                    .setPVertexAttributeDescriptions(VertexInputAttributeDescriptions.Data());
                
                vk::PipelineInputAssemblyStateCreateInfo PipelineInputAssemblyStateCreateInfo = vk::PipelineInputAssemblyStateCreateInfo()
                    .setTopology(vk::PrimitiveTopology::eTriangleList);

                vk::PipelineViewportStateCreateInfo PipelineViewportStateCreateInfo = vk::PipelineViewportStateCreateInfo()
                    .setViewportCount(1)
                    .setScissorCount(1);

                vk::PipelineRasterizationStateCreateInfo PipelineRasterizationStateCreateInfo = vk::PipelineRasterizationStateCreateInfo()
                    .setPolygonMode(vk::PolygonMode::eFill)
                    .setFrontFace(vk::FrontFace::eCounterClockwise)
                    .setCullMode(PreRenderPipeline.bDoubleSided? vk::CullModeFlagBits::eNone : vk::CullModeFlagBits::eBack)
                    .setLineWidth(1.0f);

                vk::PipelineMultisampleStateCreateInfo PipelineMultisampleStateCreateInfo = vk::PipelineMultisampleStateCreateInfo()
                    .setRasterizationSamples(vk::SampleCountFlagBits::e1);

                vk::PipelineDepthStencilStateCreateInfo PipelineDepthStencilStateCreateInfo = vk::PipelineDepthStencilStateCreateInfo()
                    .setDepthTestEnable(true)
                    .setDepthWriteEnable(true)
                    .setDepthCompareOp(vk::CompareOp::eLess);

                TVector<vk::PipelineColorBlendAttachmentState> PipelineColorBlendAttachmentStates;

                switch (PreRenderPipeline.Type)
                {
                    case ERenderMaterialType::Opaque:
                    {
                        vk::PipelineColorBlendAttachmentState PipelineColorBlendAttachmentState = vk::PipelineColorBlendAttachmentState()
                            .setColorWriteMask(vk::ColorComponentFlagBits::eR | vk::ColorComponentFlagBits::eG | vk::ColorComponentFlagBits::eB | vk::ColorComponentFlagBits::eA)
                            .setBlendEnable(false);

                        PipelineColorBlendAttachmentStates={  //输出合成阶段需要的图片
                            PipelineColorBlendAttachmentState
                            ,PipelineColorBlendAttachmentState
                            ,PipelineColorBlendAttachmentState
                            ,PipelineColorBlendAttachmentState
                        };
                        break;
                    }
                    case ERenderMaterialType::Unlit:
                    {
                        vk::PipelineColorBlendAttachmentState PipelineColorBlendAttachmentState = vk::PipelineColorBlendAttachmentState()
                            .setColorWriteMask(vk::ColorComponentFlagBits::eR | vk::ColorComponentFlagBits::eG | vk::ColorComponentFlagBits::eB | vk::ColorComponentFlagBits::eA)
                            .setBlendEnable(true)
                            .setSrcColorBlendFactor(vk::BlendFactor::eSrcAlpha)
                            .setDstColorBlendFactor(vk::BlendFactor::eOneMinusSrcAlpha)
                            .setColorBlendOp(vk::BlendOp::eAdd)
                            .setSrcAlphaBlendFactor(vk::BlendFactor::eZero)
                            .setDstAlphaBlendFactor(vk::BlendFactor::eOne)
                            .setAlphaBlendOp(vk::BlendOp::eAdd);

                        PipelineColorBlendAttachmentStates={  //输出合成阶段需要的图片
                            PipelineColorBlendAttachmentState
                        };
                        break;
                    }
                };

                vk::PipelineColorBlendStateCreateInfo PipelineColorBlendStateCreateInfo = vk::PipelineColorBlendStateCreateInfo()
                    .setAttachments(PipelineColorBlendAttachmentStates);

                int32_t SubpassIndex=0;
                switch(PreRenderPipeline.Type)
                {
                    case ERenderMaterialType::Opaque:
                    {
                        SubpassIndex=0;
                        break;
                    }

                    case ERenderMaterialType::Unlit:
                    {
                        SubpassIndex=2;
                        break;
                    }
                }

                vk::GraphicsPipelineCreateInfo GraphicsPipelineCreateInfo = vk::GraphicsPipelineCreateInfo()
                    .setStageCount(2)
                    .setPStages(PipelineShaderStageCreateInfos)
                    .setPVertexInputState(&PipelineVertexInputStateCreateInfo)
                    .setPInputAssemblyState(&PipelineInputAssemblyStateCreateInfo)
                    .setPViewportState(&PipelineViewportStateCreateInfo)
                    .setPRasterizationState(&PipelineRasterizationStateCreateInfo)
                    .setPMultisampleState(&PipelineMultisampleStateCreateInfo)
                    .setPDepthStencilState(&PipelineDepthStencilStateCreateInfo)
                    .setPColorBlendState(&PipelineColorBlendStateCreateInfo)
                    .setPDynamicState(&PipelineDynamicStateCreateInfo)
                    .setLayout(NewPipeline->PipelineLayout)
                    .setRenderPass(CVulkanContext::Get()->GetSceneGlobal()->SceneRenderPass)
                    .setSubpass(SubpassIndex);

                NewPipeline->Pipeline = CVulkanContext::Get()->GetDevice().createGraphicsPipeline(nullptr, GraphicsPipelineCreateInfo).value;
            }    
            
        }

        auto Pipeline = MaterialPipelines[PreRenderPipeline];

        auto NewMaterial = std::make_shared<CVulkanMaterial>();
        Materials.Add(MaterialData, NewMaterial);

        NewMaterial->Parameters = Parameters;
        NewMaterial->Pipeline = Pipeline;

        // //创建描述符集，由于描述符池自动重置，所以每帧都要重新分配描述符
        // {
        //     TVector<vk::DescriptorSetLayout> DescriptorSetLayouts = { Pipeline->DescriptorSetLayout };

        //     vk::DescriptorSetAllocateInfo DescriptorSetAllocateInfo = vk::DescriptorSetAllocateInfo()
        //     .setDescriptorPool(CVulkanContext::Get()->GetAutoResetDescriptorPool()->MaterialPool.Pool)
        //     .setSetLayouts(DescriptorSetLayouts);            

        //     NewMaterial->DescriptorSet = CVulkanContext::Get()->GetDevice().allocateDescriptorSets(DescriptorSetAllocateInfo)[0];
        // }

        // //写入描述符
        // {
        //     TVector<vk::WriteDescriptorSet> WriteDescriptorSets;

        //     int32_t CurrentBindingIndex = 0;
        //     TVector<vk::DescriptorImageInfo> DescriptorTexture2DInfos;
        //     if(!Parameters->Texture2Ds.Empty())
        //     {
        //         for (int32_t i = 0; i < Parameters->Texture2Ds.Num(); i++)
        //         {
        //             DescriptorTexture2DInfos.Add(vk::DescriptorImageInfo()
        //                 .setImageLayout(vk::ImageLayout::eShaderReadOnlyOptimal)
        //                 .setImageView(Parameters->Texture2Ds[i]->ImageView)
        //                 .setSampler(CVulkanContext::Get()->GetDefaultImageSampler()));
        //         }

        //         WriteDescriptorSets.Add(vk::WriteDescriptorSet()
        //             .setDstSet(NewMaterial->DescriptorSet)
        //             .setDstBinding(CurrentBindingIndex)
        //             .setDescriptorType(vk::DescriptorType::eCombinedImageSampler)
        //             .setDescriptorCount(DescriptorTexture2DInfos.Num())
        //             .setPImageInfo(DescriptorTexture2DInfos.Data())
        //             .setDstArrayElement(0));
                
        //         CurrentBindingIndex++;         
        //     }

        //     TVector<vk::DescriptorImageInfo> DescriptorTextureCubeImageInfos;
        //     if(!Parameters->TextureCubes.Empty())
        //     {
        //         for (int32_t i = 0; i < Parameters->TextureCubes.Num(); i++)
        //         {
        //             DescriptorTextureCubeImageInfos.Add(vk::DescriptorImageInfo()
        //                 .setImageLayout(vk::ImageLayout::eShaderReadOnlyOptimal)
        //                 .setImageView(Parameters->TextureCubes[i]->ImageView)
        //                 .setSampler(Parameters->TextureCubes[i]->Sampler));
        //         }

        //         WriteDescriptorSets.Add(vk::WriteDescriptorSet()
        //             .setDstSet(NewMaterial->DescriptorSet)
        //             .setDstBinding(CurrentBindingIndex)
        //             .setDescriptorType(vk::DescriptorType::eCombinedImageSampler)
        //             .setDescriptorCount(DescriptorTextureCubeImageInfos.Num())
        //             .setPImageInfo(DescriptorTextureCubeImageInfos.Data())
        //             .setDstArrayElement(0));
                
        //         CurrentBindingIndex++;         
        //     }

        //     vk::DescriptorBufferInfo DescriptorBufferInfo = vk::DescriptorBufferInfo()
        //     .setBuffer(Parameters->FloatParameterBuffer.GetBuffer())
        //     .setOffset(0)
        //     .setRange(Parameters->FloatParameterBuffer.GetSize());
        //     if(Parameters->FloatParameterBuffer.Size()>0)
        //     {

        //         WriteDescriptorSets.Add(vk::WriteDescriptorSet()
        //             .setDstSet(NewMaterial->DescriptorSet)
        //             .setDstBinding(CurrentBindingIndex)
        //             .setDescriptorType(vk::DescriptorType::eUniformBuffer)
        //             .setDescriptorCount(1)
        //             .setPBufferInfo(&DescriptorBufferInfo)
        //             .setDstArrayElement(0));
                
        //         CurrentBindingIndex++;
        //     }

        //     CVulkanContext::Get()->GetDevice().updateDescriptorSets(WriteDescriptorSets, nullptr);

        // }
    }

}

std::shared_ptr<CVulkanMaterial> CVulkanMaterials::FindVulkanObj(std::shared_ptr<CPreRenderMaterial> InMaterial)
{
    return Materials[InMaterial];
}
